#!  /usr/bin/env python
# coding=utf-8

import rospy
import socket
import os
import json
import numpy as np
import math
import cv2 as cv
from ti5arm.msg import ti5_polyJC_Lists
from ti5arm.msg import ti5_motorStatus

pi = math.pi

from ti5_fkik_24 import TFMeef, ti5_optIK, ti5FK



def writePToFile(polyPoseLists, cmd_rate, usetime, step, dir='None'):
    # type:(np.ndarray, float, float, int, str)->None
    row,col = polyPoseLists.shape
    if dir=='None':
        dir = '/home/wheeltec/Desktop/ti5Poly_writePLists.txt'
    
    with open(dir, 'w') as f:
        head = {'name':'polyPoseLists', 'writeTime':rospy.Time.now().to_sec(), 
                 'cmdRate':cmd_rate, 'cmdTime':usetime, 'step':step}
        msg_head = json.dumps(head)
        f.write(msg_head)

        for i in range(0, row):
            w_msg = "\r\n{}, {}, {}, {}, {}, {}".format(polyPoseLists[i,0],
                                                        polyPoseLists[i,1],
                                                        polyPoseLists[i,2],
                                                        polyPoseLists[i,3],
                                                        polyPoseLists[i,4],
                                                        polyPoseLists[i,5])
            f.write(w_msg)
            pass 
        f.write("\r\n")
        print("write polyPoseLists to {}".format(dir))
        pass
    
    
    pass

def linerPoly(xyz_eulZYX, tList, step, save_jdeg=0):
    # type:(np.ndarray, np.ndarray, int, int)->np.ndarray
    step = int(step)
    poly_rate = step/(tList[-1] - tList[0])
    
    poly_stepValue = (xyz_eulZYX[-1] - xyz_eulZYX[0])/step
   
    polyJdegLists = np.zeros((step+1,6), dtype=np.float)
    
    poly_poseLists = np.zeros((step+1,6), dtype=np.float)

    T = TFMeef(xyz_eulZYX[0])
    polyJdegLists[0] = (ti5_optIK(cJdeg=np.zeros((1,6),np.float),Targ=T))
    poly_poseLists[0] = xyz_eulZYX[0,:] 
    print('poly_poseLists[0]')
    print(poly_poseLists[0])

    for i in range(1,step+1):
        poly_poseLists[i] = xyz_eulZYX[0] + poly_stepValue*i
        T = TFMeef(poly_poseLists[i])
        polyJdegLists[i] = ti5_optIK(cJdeg=polyJdegLists[i-1],Targ=T)
        
        pass
    print('poly_poseLists[-1]')
    print(poly_poseLists[-1])
    if save_jdeg:
        writePToFile(poly_poseLists, poly_rate, tList[-1] - tList[0], step)

    return polyJdegLists
    pass

def cartesian_LP_ms(xyz_eulZYX, tList, cmdRate, cmdTime, save_jdeg=0):
    # type:(np.ndarray, np.ndarray, int, float, int)-> np.ndarray
    """cartesian linear Poly. Use  cmdRate and stepValue move cmdTime(s) along the point specified in xyz_eulZYX
       
        input:
            xyz_eulZYX: cartesian pose/point list (xyz and eul_ZYX)
            tList: time corresponding to pose
        output:
            polyJdegLists: interpolated joint angle(degree)
            
    """
    step = int(cmdRate*cmdTime)
    poly_rate = cmdRate
    endPose = xyz_eulZYX[-1]
    pointDist = xyz_eulZYX[-1] - xyz_eulZYX[0]
    ratio = cmdTime/abs(tList[-1]-tList[0])
    for i in range(0,2): # just poly x y
        endPose[i] =xyz_eulZYX[0][i] + pointDist[i]*ratio
        pass
    poly_stepValue = (endPose - xyz_eulZYX[0])/step
   
    polyJdegLists = np.zeros((step+1,6), dtype=np.float)
    
    poly_poseLists = np.zeros((step+1,6), dtype=np.float)

    T = TFMeef(xyz_eulZYX[0])
    polyJdegLists[0] = (ti5_optIK(cJdeg=np.zeros((1,6),np.float),Targ=T))
    poly_poseLists[0] = xyz_eulZYX[0,:] 
   
    print('poly_poseLists[0]',poly_poseLists[0])
    for i in range(1,step+1):
        poly_poseLists[i] = xyz_eulZYX[0] + poly_stepValue*i
        T = TFMeef(poly_poseLists[i])
        polyJdegLists[i] = ti5_optIK(cJdeg=polyJdegLists[i-1],Targ=T)
        
        pass
    print('endPose',endPose)
    if save_jdeg:
        writePToFile(poly_poseLists, poly_rate, cmdTime, step)

    return polyJdegLists
    pass

def pub_jLists(pub,jLists,cmd_rate, time):
    # # type:(Any,np.ndarray,int,float)->None

    # 如果需要自动运行令auto_run=1
    auto_run = 1 # 默认不运行
    
    jointTarget = ti5_polyJC_Lists()
    # **************************config ctrl mode********************************
    jointTarget.numOfCmd = 7
    jointTarget.interpolyReq = 0 # not need interpoly
    jointTarget.mode = 0
    jointTarget.header.stamp = rospy.Time().now()
    jointTarget.cmdRate = cmd_rate
    # **************************************************************************

    print("*********+++++pub_jLists start+++*********")
    if auto_run==1:
        #jointxRadList is deg
        
        nJ = np.zeros((int(np.size(jLists)/6),7), np.float)
        nJ[:,6] = -600000
        nJ[-int((np.size(jLists)*2)/15):-1,6] = 0 # last 80% point will close gripper
        nJ[-1,6] = 0
        nJ[:,0:6] = jLists
        
        print('***')
        print(jLists.shape[0])
        for i in range(0,nJ.shape[0]):
            jointTarget.joint1RadList.append(nJ[i,0])
            jointTarget.joint2RadList.append(nJ[i,1])
            jointTarget.joint3RadList.append(nJ[i,2])
            jointTarget.joint4RadList.append(nJ[i,3])
            jointTarget.joint5RadList.append(nJ[i,4])
            jointTarget.joint6RadList.append(nJ[i,5])
            jointTarget.joint7RadList.append(nJ[i,6])
        # print('jointTarget:')
        # for i in range(0,len(jointTarget.joint1RadList)):
        #     print(jointTarget.joint1RadList[i],\
        #           jointTarget.joint2RadList[i],\
        #           jointTarget.joint3RadList[i],\
        #           jointTarget.joint4RadList[i],\
        #           jointTarget.joint5RadList[i],\
        #           jointTarget.joint6RadList[i],\
        #           jointTarget.joint7RadList[i])

        pub.publish(jointTarget)
        print("pub_jLists ok")
       
    pass

def testMotor_LinerPoly(sJDeg, eJDeg, pub,cmd_rate,use_time):
    # 如果需要自动运行令auto_run=1
    auto_run = 1 # 
    
    jointTarget = ti5_polyJC_Lists()
    # **************************config ctrl mode********************************
    jointTarget.numOfCmd = 6
    jointTarget.interpolyReq = 10 # iner linear interpoly
    jointTarget.mode = 0
    jointTarget.header.stamp = rospy.Time().now()
    jointTarget.cmdRate = int(cmd_rate)
    # **************************************************************************
    # print("*********+++++LinerPoly start+++*********")
    
    # from pose 0 to pose 1
    if auto_run==1:
        #jointxRadList is deg
      
        nJ = [sJDeg[0],    sJDeg[1],  sJDeg[2],   sJDeg[3],    sJDeg[4],    sJDeg[5],  0.0,
              eJDeg[0],    eJDeg[1],  eJDeg[2],   eJDeg[3],    eJDeg[4],    eJDeg[5],  0.0]

           
        for i in range(0,len(nJ)/7):
            jointTarget.joint1RadList.append(nJ[0+i*7])
            jointTarget.joint2RadList.append(nJ[1+i*7])
            jointTarget.joint3RadList.append(nJ[2+i*7])
            jointTarget.joint4RadList.append(nJ[3+i*7])
            jointTarget.joint5RadList.append(nJ[4+i*7])
            jointTarget.joint6RadList.append(nJ[5+i*7])
            jointTarget.joint7RadList.append(nJ[6+i*7])

        jointTarget.timeLists = [0, use_time]
    
        pub.publish(jointTarget)
        print("LinerPoly pub ok")
   
    pass
 
def JLinerPoly(sJDeg, eJDeg, pub,cmd_rate,use_time):
    
    jointTarget = ti5_polyJC_Lists()
    # **************************config ctrl mode********************************
    jointTarget.numOfCmd = 6
    jointTarget.interpolyReq = 0 # external J space linear interpoly
    jointTarget.mode = 0
    jointTarget.header.stamp = rospy.Time().now()
    jointTarget.cmdRate = int(cmd_rate)
    # **************************************************************************
    
    poly_step = int(use_time*cmd_rate)
    #jointxRadList is deg
    sJ = np.array(sJDeg, np.float)
    eJ = np.array(eJDeg, np.float)
    cnJ = np.array(eJ-sJ,np.float)
    stepValue = np.array(cnJ/poly_step,np.float)

    for i in range(1,poly_step+1):
        jointTarget.joint1RadList.append(sJDeg[0] + (i*stepValue[0]))
        jointTarget.joint2RadList.append(sJDeg[1] + (i*stepValue[1]))
        jointTarget.joint3RadList.append(sJDeg[2] + (i*stepValue[2]))
        jointTarget.joint4RadList.append(sJDeg[3] + (i*stepValue[3]))
        jointTarget.joint5RadList.append(sJDeg[4] + (i*stepValue[4]))
        jointTarget.joint6RadList.append(sJDeg[5] + (i*stepValue[5]))
   
    pub.publish(jointTarget)
    print("LinerPoly pub ok")
    
    pass

def writeJToFile(polyJdegLists, cmd_rate, usetime, step, dir = '/home/wheeltec/Desktop/ti5IK_writeJdegList.txt'):
    # type:(np.ndarray, float, float, int, str)->None
    row,col = polyJdegLists.shape
    # dir = '/home/wheeltec/Desktop/ti5IK_writeJdegList.txt'
    # try:
    with open(dir, 'w') as f:
        head = {'name':'JdegLists', 'writeTime':rospy.Time.now().to_sec(), 
                 'cmdRate':cmd_rate, 'cmdTime':usetime, 'step':step}
        msg_head = json.dumps(head)
        f.write(msg_head)
    
        for i in range(0, row):
            w_msg = "\r\n{}, {}, {}, {}, {}, {}".format(polyJdegLists[i,0],
                                                                                polyJdegLists[i,1],
                                                                                polyJdegLists[i,2],
                                                                                polyJdegLists[i,3],
                                                                                polyJdegLists[i,4],
                                                                                polyJdegLists[i,5])
            f.write(w_msg)
            pass
        f.write("\r\n")
        print("write polyJdegLists to {}".format(dir))
        pass
    # except Exception as e:
    #     print('write polyJdegLists to file failed')
    #     print(e)


    pass

def sub_motorStat():
    
    """堵塞20ms至收到关节消息/deg,返回的关节角为二维数组
    """
    timeout = rospy.Duration(0,200000000) # 200ms
    try:
        motorStat = rospy.wait_for_message("ti5_motor_status", ti5_motorStatus, timeout=timeout) # type:ti5_motorStatus
    except Exception as e:
        print(e)
        return -1
        

    try:
        timeStamp = motorStat.header.stamp
        numOfJ = motorStat.numOfMotor

        jLists = np.array(motorStat.joint1_7DegLists, dtype=np.float) # type:np.ndarray

        step = int(jLists.size/numOfJ)
        jLists.resize((step,numOfJ))
        if jLists.size >0:
            # print('jLists')
            # print(jLists)
            # print(jLists.shape)
            
            return jLists
            pass
        else:
            print('jLists.size',jLists.size)
            return -1
        pass
    except Exception as e:
        print('can not sub msg: "ti5_motor_status" ')
        return -1
    
    pass

def get_msg_pubCMD(connect, pub, save_polyJdeg=0):
    s_pose_jdeg7 = np.array([0,0,-90,0,-90,-90,-32.6],dtype=np.float)
    s_pose_eulZYX_deg3 = np.array([-90,0.0,180.0],dtype=np.float)
    z = 70 # safety distance limit while gripper fully closed(/mm)
    grabTime = 0.5 # /s
    pose_flag = 0 # not ready
    cmd_flag = 0 # motor stoped
    tE = 0.5 # time error

    # ********24/2/25 modified********
    X_err = 20 # 坐标轴修正误差，由标定误差决定(/mm)
    Y_err = 0
    Z_err = 0
    # *******************************
    arm_cmdRate = 500

    pose_flag = 0
    motorStat = sub_motorStat() # get motor sta

    imge = np.zeros((480,460,3), dtype=np.uint8)#要显示的UI图像
    cv.imshow('ti5GrabFruits', imge)
    while True:
        key = cv.waitKey(2)
        if key & 0xff == ord('q'):
            cv.destroyAllWindows()
            break

        if cmd_flag == 1 or pose_flag < 2 : # motor running or pose not ready
        
            motorStat = sub_motorStat() # get motor sta
            pass
        

        recv = connect.recv(1000).decode() # get msg

        if cmd_flag==1 and type(-1)==type(motorStat): # motor running and get nothing
            # pose_flag = 0 # not ready

            connect.send('motor running'.encode())
            print('motor running')
            
            continue
            pass
        if  pose_flag < 2 and type(-1)!=type(motorStat): # pose not ready and get motor sta
            if (np.sum(abs(motorStat[0] - s_pose_jdeg7))>0.7): # motors are not ready sta
                pose_flag = 0 # not ready
                print('current state is not ready')
                connect.send('current state is not ready'.encode())
                print('dist between ready pose and curr Pose',np.sum(abs(motorStat[0] - np.array([0,0,-90,0,0,0,-32.6], np.float))))
                print(motorStat[0])
                continue
            else: # mortors are at ready pose
                pose_flag = pose_flag+1 # pose ready
                if pose_flag>=2:
                    cmd_flag = 0 # motor stoped
                    pass
                pass
                print('dist between ready pose and curr Pose',np.sum(abs(motorStat[0] - np.array([0,0,-90,0,0,0,-32.6], np.float))))
                print(motorStat[0])

            motorStat = -1
            pass
        
        
        if recv:
            if recv=="ok":
                # print("ooo")
                connect.send('ooo'.encode())
                continue
            else: # recved msg
                print('recv')
                print(recv)
            if pose_flag<2: # pose not ready
                print('receive a msg, but current state is not ready, msg was ignored!')
                connect.send('current state is not ready, previous msg was ignored!'.encode())
                continue
                
            data = json.loads(recv)
            print('data')
            print(type(data))
            print(data)
            x1 = data['X1']*1000 + X_err # 由于相机标定误差
            x2 = data['X2']*1000 + X_err # 位置参数需要修正
            y1 = data['Y1']*1000 + Y_err
            y2 = data['Y2']*1000 + Y_err
            z1 = data['Z1']*1000 + Z_err
            z2 = data['Z2']*1000 + Z_err
            t1 = data['T1']
            t2 = data['T2']

            T1 = TFMeef(np.array([x1,y1,z,s_pose_eulZYX_deg3[0]*pi/180.0,s_pose_eulZYX_deg3[1]*pi/180.0,s_pose_eulZYX_deg3[2]*pi/180.0],np.float))
            print('TfM1')
            print(T1)

            jdeg1 = ti5_optIK(cJdeg=s_pose_jdeg7[0:6],Targ=T1)
            print('j1')
            print(np.around(jdeg1,3))

            T2 = TFMeef(np.array([x2,y2,z,s_pose_eulZYX_deg3[0]*pi/180.0,s_pose_eulZYX_deg3[1]*pi/180.0,s_pose_eulZYX_deg3[2]*pi/180.0],np.float))
            print('TfM2')
            print(T2)

            jdeg2 = ti5_optIK(cJdeg=np.array(jdeg1,np.float),Targ=T2)
            print('j2')
            print(np.around(jdeg2,3))


            distance = sum(abs(jdeg2 - jdeg1))
            if 0.5<distance: # msg ok
                print(rospy.Time.now().to_nsec())
                time1=rospy.Time.now().to_nsec()
                openGripper(pub, -600000.0) # open gripper
                testMotor_LinerPoly(s_pose_jdeg7[0:6] ,jdeg1, pub, 200, t1-tE) # move to point1
                print('move to point1')

                cmd_rate = arm_cmdRate

                # # follow more
                # polyJdegLists = linerPoly(np.array([[x1,y1,z,90*pi/180,0,180*pi/180],
                #                                     [x2,y2,z,90*pi/180,0,180*pi/180]],dtype=np.float), 
                #                           np.array([0,t2],dtype=np.float), 
                #                           cmd_rate*t2)

                # follow less
                polyJdegLists = cartesian_LP_ms(np.array([[x1,y1,z,s_pose_eulZYX_deg3[0]*pi/180.0,s_pose_eulZYX_deg3[1]*pi/180.0,s_pose_eulZYX_deg3[2]*pi/180.0],
                                                          [x2,y2,z,s_pose_eulZYX_deg3[0]*pi/180.0,s_pose_eulZYX_deg3[1]*pi/180.0,s_pose_eulZYX_deg3[2]*pi/180.0]],dtype=np.float), 
                                                np.array([0,t2],dtype=np.float), 
                                                cmd_rate,
                                                grabTime)

                print('polyJdegLists.shape', polyJdegLists.shape)
                
                # # follow more
                # pub_jLists(pub, polyJdegLists,cmd_rate, t2)
                
                # follow less
                pub_jLists(pub, polyJdegLists,cmd_rate, grabTime) # move toward point2 and grab


                TGup = ti5FK(polyJdegLists[-1]*(pi/180))
                TGup[2,3] = 200
                print('TGup')
                print(TGup)

                jdegGup = ti5_optIK(cJdeg=polyJdegLists[-1],Targ=TGup)
                print('jdegGup')
                print(np.around(jdegGup,3))

                TPutUp = TFMeef(np.array([0,-450,200,s_pose_eulZYX_deg3[0]*pi/180.0,s_pose_eulZYX_deg3[1]*pi/180.0,s_pose_eulZYX_deg3[2]*pi/180.0],np.float))
                print('TPutUp')
                print(TPutUp)

                jdegPutUp = ti5_optIK(cJdeg=np.array(jdegGup,np.float),Targ=TPutUp)
                print('jdegPutUp')
                print(np.around(jdegPutUp,3))

                TPutDown = TFMeef(np.array([0,-450,-200,s_pose_eulZYX_deg3[0]*pi/180.0,s_pose_eulZYX_deg3[1]*pi/180.0,s_pose_eulZYX_deg3[2]*pi/180.0],np.float))
                print('TPutDown')
                print(TPutDown)

                jdegPutDown = ti5_optIK(cJdeg=np.array(jdegPutUp,np.float),Targ=TPutDown)
                print('jdegPutDown')
                print(np.around(jdegPutDown,3))
                
              
                # # follow less
                # testMotor_LinerPoly(polyJdegLists[-1], [0,0,-90,0,0,0], pub, 500, 6.0) # back to ready pose

                # move to grab up pose
                testMotor_LinerPoly(polyJdegLists[-1], jdegGup, pub, 500, 1.0)

                # move to put up pose
                testMotor_LinerPoly(jdegGup, jdegPutUp, pub, 500, 3.0)
               
                # move to put down pose
                testMotor_LinerPoly(jdegPutUp, jdegPutDown, pub, 500, 3.0)

                openGripper(pub, -600000.0) # open gripper
                
                # move to put up pose
                jdegPutUp[5]=0.0
                testMotor_LinerPoly(jdegPutDown, jdegPutUp, pub, 500, 3.0)
               
                # move to ready pose
                testMotor_LinerPoly(jdegPutUp, s_pose_jdeg7[0:6], pub, 500, 2.0)
               

                print(rospy.Time.now().to_nsec())
                time2 = rospy.Time.now().to_nsec()
                print((time2-time1)/1000000)

                # write polyJdegLists to file
                # dir = '/home/wheeltec/Desktop/ti5put_JdegList.txt'
                if save_polyJdeg:
                    writeJToFile(polyJdegLists, cmd_rate, grabTime, cmd_rate*grabTime)

                connect.send('ok'.encode())
                cmd_flag = 1 # motor running
                pose_flag = 0 # pose not ready
            else:
                print('error data')
                connect.send('error data'.encode())
            pass # end if recv
        else:
            print('error data')
            pass

    pass

def sockLink(localhost, local_port, pub):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    print('connet to client...')
    s.bind((localhost, local_port))
    s.listen(1)
    print('server is listening...')
    client_sock, client_addr = s.accept()
    print('addr',client_addr)
    print('connet successfully!')
    get_msg_pubCMD(client_sock, pub)
    client_sock.close()
    print('connet closed!')
    pass

def openGripper(pub, pose):
    jointTarget = ti5_polyJC_Lists()
    # **************************config ctrl mode********************************
    jointTarget.numOfCmd = 1
    jointTarget.interpolyReq = 10 # iner linear interpoly
    jointTarget.mode = 0
    jointTarget.header.stamp = rospy.Time().now()
    jointTarget.cmdRate = int(500)
    # **************************************************************************
    print("*********+++++openGripper start+++*********")
    nJ = [  0.0,    0.0,   0.0,     0.0,    0.0,   0.0,  float(pose)]
    for i in range(0,len(nJ)/7):
        jointTarget.joint1RadList.append(nJ[0+i*7])
        jointTarget.joint2RadList.append(nJ[1+i*7])
        jointTarget.joint3RadList.append(nJ[2+i*7])
        jointTarget.joint4RadList.append(nJ[3+i*7])
        jointTarget.joint5RadList.append(nJ[4+i*7])
        jointTarget.joint6RadList.append(nJ[5+i*7])
        jointTarget.joint7RadList.append(nJ[6+i*7])

        jointTarget.timeLists = [0, 4]
    
        pub.publish(jointTarget)
        print("openGripper pub ok")


    pass

def testPolyResult(dirJ, dirP, line):
    
    # read JdegLists**********************
    with open(dirJ, mode='r') as Jfile:
        Jdata = Jfile.readlines()
        pass
    groups_jdeg = (len(Jdata)-1)
    try:
        head = json.loads(Jdata[0])
        name = head['name']
        cmdTime = head['cmdTime']
        cmdRate = head['cmdRate']
        step = head['step']
    except Exception as e:
        print('json loads msg_head failed!')
        print(e)
        name = ''

    if(name=="JdegLists"):# 
        print("open file!")
        print("groups of Jdeg:",groups_jdeg)
        Jdata_split = Jdata[line-1].split(', ')
        Jdeg = np.zeros(6, dtype=np.float)
        for i in range(0,6):
            Jdeg[i] = float(Jdata_split[i])
        print("line={}\r\nreadJdeg=\r\n{}".format(line,Jdeg))
    
    # read polyPoseLists**********************
    with open(dirP, mode='r') as Pfile:
        Pdata = Pfile.readlines()
        pass
    groups_P = (len(Pdata)-1)

    try:
        head = json.loads(Pdata[0])
        name = head['name']
        cmdTime = head['cmdTime']
        cmdRate = head['cmdRate']
        step = head['step']
    except Exception as e:
        print('json loads msg_head failed!')
        print(e)
        name = ''
    if(name=="polyPoseLists"):# 
        print("open file!")
        print("groups of Pose:",groups_P)
        Pdata_split = Pdata[line-1].split(', ')
        Pose = np.zeros(6, dtype=np.float)
        for i in range(0,6):
            Pose[i] = float(Pdata_split[i])
     
        print("line={}\r\nreadPose=\r\n{}".format(line,Pose))

    

    T_poly_pose = TFMeef(Pose)
    print("TFMeef M by read Pose:\r\n{}".format(T_poly_pose))

    IK_deg = ti5_optIK(cJdeg=np.zeros((1,6),np.float),Targ=T_poly_pose)
    print("IK deg by read Pose:\r\n{}".format(IK_deg))

    FKxyz=ti5FK(Jdeg*(pi/180))
    print("FK M by read Jdeg:\r\n{}".format(FKxyz))

    FKik=ti5FK(IK_deg*(pi/180))
    print("FK M by IK_deg:\r\n{}".format(FKik))
  

    pass



if  __name__== "__main__":
 
    rospy.init_node("ti5GrabFruits_node")

    #**********连接视觉，开始抓取***********
    pub = rospy.Publisher("interPolyPoint_topic",ti5_polyJC_Lists,queue_size=1000)

    rate = rospy.Rate(1)
    rate.sleep()
    rate.sleep()
    rate.sleep()#waiter for system

    localhost = '0.0.0.0'
    local_port = 38899

    openGripper(pub, -600000.0)
    sockLink(localhost, local_port, pub)
    #************************************
 

    pass

 





 


  

   

